import torch
import torch.nn as nn
import torch.nn.functional as F

'''
input A


resNet
B = g(A) +f(A)


B1 = f1(A)
B2 = f2(A)
B3 = f3(A)
concat([B1,B2,B3])
'''


def ConvBNRelu(in_channel, out_channel, kernel_size):
    return nn.Sequential(
        nn.Conv2d(in_channel, out_channel, kernel_size=kernel_size, stride=1, padding=kernel_size // 2),
        nn.BatchNorm2d(out_channel),
        nn.ReLU()
    )

# BaseInception(in_channel=128, out_channel_list=[64, 64, 64, 64], reduce_channel_list=[16, 16]),
class BaseInception(nn.Module):
    def __init__(self, in_channel, out_channel_list, reduce_channel_list):
        super(BaseInception, self).__init__()
        self.branch1_conv = ConvBNRelu(in_channel, out_channel_list[0], 1)

        self.branch2_conv1 = ConvBNRelu(in_channel, reduce_channel_list[0], 1)
        self.branch2_conv2 = ConvBNRelu(reduce_channel_list[0], out_channel_list[1], 3)

        self.branch3_conv1 = ConvBNRelu(in_channel, reduce_channel_list[1], 1)
        self.branch3_conv2 = ConvBNRelu(reduce_channel_list[1], out_channel_list[2], 5)

        self.branch4_pool = nn.MaxPool2d(kernel_size=3, stride=1, padding=1)
        self.branch4_conv = ConvBNRelu(in_channel, out_channel_list[3], 3)

    def forward(self, x):
        out_1 = self.branch1_conv(x)
        # torch.Size([128, 64, 14, 14])
        print("out_1 shape:", out_1.shape)
        out_2 = self.branch2_conv1(x)
        # torch.Size([128, 16, 14, 14])
        print("out_2 shape:", out_2.shape)
        out_2 = self.branch2_conv2(out_2)
        # torch.Size([128, 64, 14, 14])
        print("out_2 shape:", out_2.shape)

        out_3 = self.branch3_conv1(x)
        # torch.Size([128, 16, 14, 14])
        print("out_3 shape:", out_3.shape)
        out_3 = self.branch3_conv2(out_3)
        # torch.Size([128, 64, 14, 14])
        print("out_3 shape:", out_3.shape)
        out_4 = self.branch4_pool(x)
        # torch.Size([128, 128, 14, 14])
        print("out_4 shape:", out_4.shape)
        out_4 = self.branch4_conv(out_4)
        # torch.Size([128, 64, 14, 14])
        print("out_4 shape:", out_4.shape)

        # batch_size不变    channel累计合并
        out5 = torch.cat([out_1, out_2, out_3, out_4], dim=1)
        # torch.Size([128, 256, 14, 14])
        print("out5 shape:", out5.shape)
        return out5



class InceptionNet(nn.Module):
    def __init__(self):
        super(InceptionNet, self).__init__()
        # =====================================================================
        # 卷积层 3*28*28 -》 64*14*14
        self.block1 = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU()
        )
        self.block2 = nn.Sequential(
            nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU()
        )

        self.block3 = nn.Sequential(
            BaseInception(in_channel=128, out_channel_list=[64, 64, 64, 64], reduce_channel_list=[16, 16]),
            nn.MaxPool2d(kernel_size=3, stride=2, padding=1))

        self.block4 = nn.Sequential(
            BaseInception(in_channel=256, out_channel_list=[96, 96, 96, 96], reduce_channel_list=[32, 32]),
            nn.MaxPool2d(kernel_size=3, stride=2, padding=1))

        self.fc = nn.Linear(384, 10)

    def forward(self, x):
        # x shape: torch.Size([128, 3, 28, 28])
        print("x shape:", x.shape)
        out = self.block1(x)
        # block1 shape: torch.Size([128, 64, 28, 28])
        print("block1 shape:", out.shape)
        out = self.block2(out)
        # block2 shape: torch.Size([128, 128, 14, 14])
        print("block2 shape:", out.shape)
        out = self.block3(out)
        # block3 shape: torch.Size([128, 256, 7, 7])
        print("block3 shape:", out.shape)
        out = self.block4(out)
        # block4 shape: torch.Size([128, 384, 4, 4])
        print("block4 shape:", out.shape)
        out = F.avg_pool2d(out, 4)
        # avg_pool2d shape: torch.Size([128, 384, 1, 1])
        print("avg_pool2d shape:", out.shape)
        out = out.view(out.size(0), -1)
        out = self.fc(out)
        return out


def inceptionnet():
    return InceptionNet()
